﻿using System;
using System.Runtime.InteropServices;
using Ast.Core.Services.Extensions;
using Ast.Core.WinApi;
using Ast.Framework;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface IPsLoadedModuleListService : IDependency
    {
        void ListModules(Func<PsLoadedModuleListService.LdrDataTableEntry, string, string, bool> procModuleFunc);
    }

    public class PsLoadedModuleListService : IPsLoadedModuleListService
    {
        private readonly IDeviceController _deviceController;
        private readonly IKernelAddressFinder _kernelAddressFinder;

        public PsLoadedModuleListService(IDeviceController deviceController, IKernelAddressFinder kernelAddressFinder)
        {
            _deviceController = deviceController;
            _kernelAddressFinder = kernelAddressFinder;

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public void ListModules(Func<LdrDataTableEntry, string, string, bool> procModuleFunc)
        {
            var psLoadedModuleList = _kernelAddressFinder.FindPsLoadedModuleList();
            var listHead = _deviceController.ReadAsIntPtr(psLoadedModuleList).Value;
            var current = _deviceController.ReadAs<LdrDataTableEntry>(listHead);
            while (current.InLoadOrderLinks.Flink != listHead)
            {
                var baseDllName = _deviceController.ReadAsUnicodeString(current.BaseDllName);
                var fullDllName = _deviceController.ReadAsUnicodeString(current.FullDllName);

                var goOn = procModuleFunc(current, fullDllName, baseDllName);
                if (!goOn) break;

                current = _deviceController.ReadAs<LdrDataTableEntry>(current.InLoadOrderLinks.Flink);
            }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        public struct LdrDataTableEntry
        {
            [MarshalAs(UnmanagedType.Struct)]
            public Structure.ListEntry InLoadOrderLinks;

            [MarshalAs(UnmanagedType.Struct)]
            public Structure.ListEntry InMemoryOrderLinks;

            [MarshalAs(UnmanagedType.Struct)]
            public Structure.ListEntry InInitializationOrderLinks;

            public IntPtr DllBase;

            public IntPtr EntryPoint;

            public uint SizeOfImage;

            public Structure.UnicodeString FullDllName;

            public Structure.UnicodeString BaseDllName; // +0x02c BaseDllName      : _UNICODE_STRING

            [MarshalAs(UnmanagedType.ByValArray, SizeConst = 0x4c)]
            private byte[] _;

            public IntPtr OriginalBase; // +0x080 OriginalBase     : Uint4B
        }
    }
}